package cn.doitedu

import org.apache.spark.sql.{Encoder, Encoders, SparkSession}
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.functions.{udaf, udf}
import org.roaringbitmap.longlong.Roaring64Bitmap

import java.io.{ByteArrayOutputStream, DataOutputStream}
import java.nio.ByteBuffer


case class Person(guid:Long,gender:String)

object Job02_TrafficAggregateUvBitmap {

  def main(args: Array[String]): Unit = {


    val spark = SparkSession.builder()
      .appName("geohash地理位置维表加工")
      .master("local")
      .enableHiveSupport()
      .getOrCreate()

    import spark.implicits._

    // 自定义函数，可以接收一组 长整数，然后返回一个bitmap的序列化字节
    val ds = spark.createDataset(Seq(Person(1, "m"), Person(2, "f"), Person(3, "m"), Person(4, "m")))

    ds.createOrReplaceTempView("tmp")

    spark.udf.register("my_bitmap_union",udaf(MyBitmapUnion))

    val bitmapCount = (bmBytes:Array[Byte])=>{
      // 反序列化
      val bm = Roaring64Bitmap.bitmapOf()
      bm.deserialize(ByteBuffer.wrap(bmBytes))

      bm.getLongCardinality
    }
    spark.udf.register("bitmap_count",udf(bitmapCount))

    spark.sql(
      """
        |
        |
        |select
        |    gender,
        |    bitmap_count(my_bitmap_union(guid) ) as user_cnt
        |from tmp
        |group by gender
        |
        |
        |""".stripMargin).show()


  }


  // select
  //    gender,
  //    函数(guid)
  // from  t
  // group by gender
  object MyBitmapUnion extends Aggregator[Long, Array[Byte], Array[Byte]] {

    // 初始化 中间聚合结构
    override def zero: Array[Byte] = {
      // 构造一个空的bitmap
      val bm = Roaring64Bitmap.bitmapOf()

      // 然后把它序列化成字节数组
      serBm(bm)

    }

    // 聚合逻辑：进来一个guid，你要怎么聚合
    override def reduce(bitmapBytes: Array[Byte], guid: Long): Array[Byte] = {

      // 先把bitmap字节反序列化成bitmap对象
      val bitmap = deSerBm(bitmapBytes)

      // 然后再往bitmap中添加guid
      bitmap.add(guid)

      // 再将添加了guid的bitmap序列化返回
      serBm(bitmap)

    }

    // 合并中间聚合结果： 就是reduce端对map端产生的多个中间聚合结果进行合并
    override def merge(bmBytes1: Array[Byte], bmBytes2: Array[Byte]): Array[Byte] = {

      // 反序列化bm1 和  bm2
      val bm1 = deSerBm(bmBytes1)
      val bm2 = deSerBm(bmBytes2)

      // 合并两个bitmap
      bm1.or(bm2)

      // 将合并好的bitmap再序列化返回
      serBm(bm1)

    }

    // 输出最终结果：
    override def finish(bmBytes: Array[Byte]): Array[Byte] = bmBytes

    override def bufferEncoder: Encoder[Array[Byte]] = Encoders.BINARY

    override def outputEncoder: Encoder[Array[Byte]] = Encoders.BINARY


    // 序列化工具
    def serBm(bitmap: Roaring64Bitmap): Array[Byte] = {
      val bout = new ByteArrayOutputStream()
      val dout = new DataOutputStream(bout)
      bitmap.serialize(dout)

      bout.toByteArray
    }

    // 反序列化工具
    def deSerBm(bmBytes: Array[Byte]): Roaring64Bitmap = {
      val bm = Roaring64Bitmap.bitmapOf()
      bm.deserialize(ByteBuffer.wrap(bmBytes))
      bm
    }
  }

}
